package main

import (
	"archive/zip"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"github.com/milvus-io/milvus-sdk-go/v2/client"
	"github.com/milvus-io/milvus-sdk-go/v2/entity"
	"github.com/streadway/amqp"
	"image"
	_ "image/png" // 导入png格式支持
	"io"
	"io/ioutil"
	"log"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

func main() {
	createMilvusCollect()

	walkFiles()
}

func walkFiles() {

	// 要遍历的文件夹路径
	//folderPath := "zip"
	//folderPath := "C:\\Users\\Administrator\\Desktop\\test-zip"
	folderPath := "D:\\midjourney-genernal1-4\\demo00"
	uploadNum := 1
	// 遍历文件夹
	err := filepath.Walk(folderPath, func(path string, info os.FileInfo, err error) error {
		// 忽略文件夹本身和非 ZIP 文件
		if info.IsDir() || !strings.HasSuffix(strings.ToLower(path), ".zip") {
			return nil
		}

		// 解压缩 ZIP 文件
		zipReader, err := zip.OpenReader(path)
		if err != nil {
			log.Printf("failed to open ZIP file %s: %v", path, err)
			return nil
		}
		defer zipReader.Close()

		// 遍历 ZIP 文件中的文件
		for _, file := range zipReader.File {
			// 处理 JSON 文件
			if strings.HasSuffix(strings.ToLower(file.Name), ".json") {
				jsonFile, err := file.Open()
				if err != nil {
					log.Printf("failed to open JSON file %s in ZIP %s: %v", file.Name, path, err)
					continue
				}
				defer jsonFile.Close()

				// 读取并解析 JSON 文件
				var jsonData map[string]string
				if err := json.NewDecoder(jsonFile).Decode(&jsonData); err != nil {
					log.Printf("failed to decode JSON file %s in ZIP %s: %v", file.Name, path, err)
					continue
				}

				// 获取 JSON 文件名（去掉后缀）
				jsonFileName := strings.TrimSuffix(file.Name, filepath.Ext(file.Name))
				// 如果 JSON 文件名以 ".meta" 结尾，则去掉 ".meta" 后缀
				jsonFileName = strings.TrimSuffix(jsonFileName, ".meta")
				// 获取 JSON 对应的 PNG 文件名
				pngFileName := fmt.Sprintf("%s.png", jsonFileName)
				// 找到对应的 PNG 文件
				for _, zipFile := range zipReader.File {
					if zipFile.Name == pngFileName {
						prompt := jsonData["prompts"]
						// 如果字符串中包含 "http"，直接返回原始字符串
						if strings.Contains(prompt, "http") {
							continue
						}

						// 打开 PNG 文件
						pngFile, err := zipFile.Open()
						// 使用 strings.Split 函数将字符串按照逗号分割成多个部分
						if err != nil {
							log.Fatalf("打开png文件失败：%v", err)
						}
						fileName := zipFile.Name
						oldSize := zipFile.FileInfo().Size()
						// 将 int64 转换为字符串
						size := strconv.FormatInt(oldSize, 10)
						uploadResp := httpUploadFile(fileName, pngFile)
						if uploadResp == nil || len(uploadResp.Resource) == 0 {
							fmt.Println("上传失败")
							continue
						}
						// 获取第一个部分并去掉末尾的空格

						pngFile2, err := zipFile.Open()
						if err != nil {
							log.Fatalf("打开png文件失败：%v", err)
						}
						imgConfig, _, err := image.DecodeConfig(pngFile2)
						if err != nil {
							//log.Fatalf("无法解码PNG文件：%v", err)
							fmt.Printf("无法解码PNG文件：%v", err)
							continue
						}
						//fmt.Println(imgConfig)
						//newUrl := uploadFile(fileName, oldSize, pngFile)
						//fmt.Println(newUrl)
						originalUrl := fmt.Sprintf("http://10.10.22.172:9092/fapi/file/%s?token=%s&file_quality=original", uploadResp.Resource, fileToken)
						//originalUrl := fmt.Sprintf("http://192.168.30.80:30700/fapi/file/%s?token=%s&file_quality=original", uploadResp.Resource, fileToken)

						galleryReq := &GetGalleryInfoRequest{
							ImageType: 1,
							MemberId:  1043,
							//MemberId:         2,
							Prompt:           prompt,
							Url:              uploadResp.Resource,
							OriginalUrl:      originalUrl,
							Size:             size,
							ContentType:      "image/png",
							Width:            int64(imgConfig.Width),
							Height:           int64(imgConfig.Height),
							FileName:         fileName,
							ModelType:        "Midjourney",
							Version:          "v 6.0",
							MidjourneyTaskId: "1774710665532256256",
							Seed:             "4913924096",
							AdditionalType:   0,
						}
						body, _ := json.Marshal(galleryReq)
						//fmt.Println(body)
						//time.Sleep(time.Second * 10)
						sendMessage(string(body))
						fmt.Printf("upload: %v\n", uploadNum)
						uploadNum++
						//sendGrpcMsg(newUrl)
						if err != nil {
							log.Printf("failed to open PNG file %s in ZIP %s: %v", pngFileName, path, err)
							continue
						}
						defer pngFile.Close()
					}
				}
			}
		}

		return nil
	})

	if err != nil {
		log.Fatal(err)
	}
}

// 创建milvus 集合
func createMilvusCollect() {
	milvusClient, err := client.NewGrpcClient(
		context.Background(), // ctx
		"10.10.22.104:19530", // addr
		//"192.168.30.80:31530", // addr
	)
	if err != nil {
		log.Fatal("failed to connect to Milvus:", err.Error())
	}

	defer milvusClient.Close()

	var (
		collectionName = "gallery2"
	)
	schema := &entity.Schema{
		CollectionName: collectionName,
		Description:    "画廊集合",
		Fields: []*entity.Field{
			{
				Name:       "id",
				DataType:   entity.FieldTypeInt64,
				PrimaryKey: true,
				AutoID:     true,
			},
			{
				Name:     "embedding",
				DataType: entity.FieldTypeFloatVector,
				TypeParams: map[string]string{
					"dim": "512",
				},
				Description: "嵌入向量",
			},
			{
				Name:        "image_type",
				DataType:    entity.FieldTypeInt32,
				Description: "图片类型，0:四宫格 1:放大2；2倍，3： 4倍放大 4: 8倍放大  5： AI无损放大 6： 降噪",
			},
			{
				Name:        "gallery_image_id",
				DataType:    entity.FieldTypeInt64,
				Description: "画廊表的id",
			},
			{
				Name:     "path",
				DataType: entity.FieldTypeVarChar,
				TypeParams: map[string]string{
					"max_length": "65535",
				},
				Description: "图片地址(可用于拼接各种分辨率)",
			},
			{
				Name:     "ai_model",
				DataType: entity.FieldTypeVarChar,
				TypeParams: map[string]string{
					"max_length": "65535",
				},
				Description: "ai模型",
			},
			{
				Name:     "model_version",
				DataType: entity.FieldTypeVarChar,
				TypeParams: map[string]string{
					"max_length": "65535",
				},
				Description: "模型版本",
			},
			{
				Name:     "prompt",
				DataType: entity.FieldTypeVarChar,
				TypeParams: map[string]string{
					"max_length": "65535",
				},
				Description: "提示词",
			},
			{
				Name:        "width",
				DataType:    entity.FieldTypeInt64,
				Description: "宽",
			},
			{
				Name:        "height",
				DataType:    entity.FieldTypeInt64,
				Description: "高",
			},
			{
				Name:        "aspectRatio",
				DataType:    entity.FieldTypeFloat,
				Description: "宽高比",
			},
			{
				Name:        "created_at",
				DataType:    entity.FieldTypeInt64,
				Description: "创建时间",
			},
			{
				Name:        "deleted_at",
				DataType:    entity.FieldTypeInt64,
				Description: "删除时间",
			},
		},
		EnableDynamicField: true,
	}

	err = milvusClient.CreateCollection(
		context.Background(), // ctx
		schema,
		2, // shardNum
	)
	if err != nil {
		log.Fatal("failed to create collection:", err.Error())
	}
	// 定义索引参数
	indexParams, _ := entity.NewIndexIvfFlat(entity.COSINE, 2048)
	// 为 embedding 字段创建索引
	err = milvusClient.CreateIndex(
		context.Background(), // ctx
		collectionName,       // collectionName
		"embedding",          // fieldName
		indexParams,          // indexParams
		true,
	)
	if err != nil {
		log.Fatal("failed to create index:", err.Error())
	}

}

// 测试版
var token = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJhdWQiOiJ4ZC1kZXNpZ24tc2VydmljZSIsImV4cCI6MTcyNzMzMDk0OSwiaWF0IjoxNzI3MDcxNzQ5LCJpc3MiOiJ4ZC1taWRqb3VybmV5LXNlcnZlciIsIklEIjoxMDQzLCJVc2VybmFtZSI6IuWwj-eBsOeBsCIsIlJvbGVJZHMiOm51bGwsIlBsYXRGb3JtIjoyLCJNZW1iZXJTdGF0dXMiOjF9.ODiPsai_QkrnXpwlpQotweoBJiY5T48P7nzPMzwEEHoxbS4FS4FSlfg9Lo_LMWIRddXWVMqLJFdRMy0sC25RKsYV7-oRh3S-SoYqj60hNfAFEUbyTWc1pAYJbrsmvCUcctB1ntZbdsWe7uDF8LD6vSqFwkk-m338X_rbf3yZUmg"
var fileToken = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJhdWQiOiJ4ZC1kZXNpZ24iLCJleHAiOjE3Mjc2NzY1NDksImlhdCI6MTcyNzA3MTc0OSwiaXNzIjoieGQtbWlkam91cm5leS1zZXJ2ZXIiLCJUb2tlblBvbGljeSI6eyIxIjpbeyJQZXJtaXNzaW9uIjoxLCJGaWxlRGlyIjoiL2FkZGl0aW9uYWwvMTA0My9vcmlnaW5hbC8ifV0sIjIiOlt7IlBlcm1pc3Npb24iOjEsIkZpbGVEaXIiOiIvcHVibGljLyovKi8ifSx7IlBlcm1pc3Npb24iOjEsIkZpbGVEaXIiOiIvKi8xMDQzLyovIn0seyJQZXJtaXNzaW9uIjoxLCJGaWxlRGlyIjoiLyovKi8qLyJ9LHsiUGVybWlzc2lvbiI6MSwiRmlsZURpciI6Ii8qLyovd2ViX2xvc3NsZXNzLyJ9LHsiUGVybWlzc2lvbiI6MSwiRmlsZURpciI6Ii8qLyovd2ViX2hpZ2gvIn0seyJQZXJtaXNzaW9uIjoxLCJGaWxlRGlyIjoiLyovKi93ZWJfbWlkZGxlLyJ9LHsiUGVybWlzc2lvbiI6MSwiRmlsZURpciI6Ii8qLyovd2ViX2xvdy8ifV19fQ.GyIXVKK1lDIIdOVB_zgQ8SOWeRY1vZt_SaQWZ1eav47ck86sJspFAMmzOYVNiCbeMNbE_vfZlBSDSI6cUJMmNKjRnKzbGH1-5XC9b_SMsp5LAl3BsL1QzLH_4HJ39RGKvXuT3AhE3ghqLRqMwcvHn08KztDYVqiNmsb1-W2r3ii1LxOhKkPaMYU7yHAonpEsEDRhA9mcQglhxnQMG_rUPe40C93YatyHv7ArwkU9s3sfAdfzzYITyld4AuOK8fnRTFgPkSqT6a7kFhuzs2YzFnTKPXNEWWa1wGDr9-4YmYmWsSiLoh1lr9wD90Qx1X2GWrymkpJjyOLZExlLeD3yHg"

// 连接 RabbitMQ 服务器
var conn, _ = amqp.Dial("amqp://admin:admin@10.10.22.172:5672/my_vhost")

var ch, _ = conn.Channel()

func sendMessage(body string) {

	err := ch.Publish(
		"aigcbus.gallery", // 交换器名称
		"gallery",         // 路由键
		false,             // 强制性
		false,             // 立即发布
		amqp.Publishing{
			ContentType: "text/plain",
			Body:        []byte(body),
		})
	failOnError(err, "Failed to publish a message")

	//fmt.Println(" [x] Sent ", body)
}
func failOnError(err error, msg string) {
	if err != nil {
		log.Fatalf("%s: %s", msg, err)
	}
}

type GetGalleryInfoRequest struct {
	ImageType        int32  `protobuf:"varint,1,opt,name=image_type,json=imageType,proto3" json:"image_type,omitempty"`
	MemberId         int64  `protobuf:"varint,2,opt,name=member_id,json=memberId,proto3" json:"member_id,omitempty"`
	Prompt           string `protobuf:"bytes,3,opt,name=prompt,proto3" json:"prompt,omitempty"`
	Url              string `protobuf:"bytes,4,opt,name=url,proto3" json:"url,omitempty"`
	OriginalUrl      string `protobuf:"bytes,5,opt,name=original_url,json=originalUrl,proto3" json:"original_url,omitempty"`
	Size             string `protobuf:"bytes,6,opt,name=size,proto3" json:"size,omitempty"`
	ContentType      string `protobuf:"bytes,7,opt,name=content_type,json=contentType,proto3" json:"content_type,omitempty"`
	Width            int64  `protobuf:"varint,8,opt,name=width,proto3" json:"width,omitempty"`
	Height           int64  `protobuf:"varint,9,opt,name=height,proto3" json:"height,omitempty"`
	FileName         string `protobuf:"bytes,10,opt,name=file_name,json=fileName,proto3" json:"file_name,omitempty"`
	ModelType        string `protobuf:"bytes,11,opt,name=model_type,json=modelType,proto3" json:"model_type,omitempty"`
	Version          string `protobuf:"bytes,12,opt,name=version,proto3" json:"version,omitempty"`
	MidjourneyTaskId string `protobuf:"bytes,13,opt,name=midjourney_task_id,json=midjourneyTaskId,proto3" json:"midjourney_task_id,omitempty"` // mj绘画id
	AdditionalType   int32  `protobuf:"varint,14,opt,name=additional_type,json=additionalType,proto3" json:"additional_type,omitempty"`        // 0 没有附加图 1是垫图 2混图
	Seed             string `protobuf:"bytes,15,opt,name=seed,proto3" json:"seed,omitempty"`                                                   //种子值
}

type UploadFileResponse struct {
	Resource string `json:"resource"`
}

// http 上传文件到AI服务器
func httpUploadFile(fileName string, file io.ReadCloser) *UploadFileResponse {

	//url := "http://192.168.30.80:30700/fapi/file"
	url := "http://10.10.22.172:9092/fapi/file"
	method := "POST"
	payload := &bytes.Buffer{}
	writer := multipart.NewWriter(payload)
	_ = writer.WriteField("file_type", "additional")
	part2,
		errFile2 := writer.CreateFormFile("file", fileName)
	_, errFile2 = io.Copy(part2, file)
	if errFile2 != nil {
		fmt.Println(errFile2)
		return nil
	}
	_ = writer.WriteField("token", fileToken)
	err := writer.Close()
	if err != nil {
		fmt.Println(err)
		return nil
	}

	client := &http.Client{}
	req, err := http.NewRequest(method, url, payload)

	if err != nil {
		fmt.Println(err)
		return nil
	}
	req.Header.Add("Content-Type", "image/png")
	req.Header.Add("Authorization", token)
	req.Header.Set("Content-Type", writer.FormDataContentType())
	res, err := client.Do(req)
	if err != nil {
		fmt.Println(err)
		return nil
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		fmt.Println(err)
		return nil
	}
	//fmt.Println(string(body))

	var response = &UploadFileResponse{}
	if err := json.Unmarshal(body, response); err != nil {
		fmt.Errorf("err=%v", err)
		return nil
	}
	return response
}
